#include <linux/blkdev.h>
#include <linux/blk-mq.h>

#include "virtual_device.h"
/*
Send IO:
	origin_request + origin_bio => clone_request + clone_bio
Complete IO:
	clone_request + clone_bio => origin_request + origin_bio

end_clone_bio(clone_bio)
  end_clone_request(clone_request, error)
    blk_mq_complete_request(origin_rq)
	  dm_softirq_done(origin_rq)

*/

static void end_clone_bio(struct bio *clone)
{
	bool is_last = !clone->bi_next;
	unsigned int nr_bytes = 0;
	blk_status_t error = clone->bi_status;
	clone_bio_t *clone_bio = NULL;
	device_request_t *dev_req = NULL;

	clone_bio = container_of(clone, clone_bio_t, clone);
	nr_bytes = clone_bio->origin->bi_iter.bi_size;
	dev_req = clone_bio->req;
	bio_put(clone);

	if (dev_req->error) {
		return;
	}
	else if (error) {
		dev_req->error = error;
		blk_update_request(dev_req->origin_req, BLK_STS_OK, dev_req->completed);
		return;
	}

	dev_req->completed += nr_bytes;
	if (is_last) {
		blk_update_request(dev_req->origin_req, dev_req->error, 0);
	}
	return;
}

static enum rq_end_io_ret end_clone_request(struct request *clone, blk_status_t error)
{
	device_request_t *dev_req = clone->end_io_data;
	dev_req->error = error;
	blk_mq_complete_request(dev_req->origin_req);
	return RQ_END_IO_FREE; /* 返回该值，内核会释放 clone request*/
}

static int rq_bio_constructor(struct bio *bio, struct bio *bio_orig, void *data)
{
	bio->bi_private = data;
	bio->bi_end_io = end_clone_bio;
	return 0;
}

static blk_status_t device_mq_queue_rq(struct blk_mq_hw_ctx *hctx, const struct blk_mq_queue_data *bd)
{
	int ret = -1;
	struct request *clone_rq = NULL;
	struct request *origin_rq = NULL;
	virtual_device_t *dev = NULL;
	device_request_t *dev_req = NULL;
	struct request_queue *path_queue = NULL;

	origin_rq = bd->rq;
	dev_req = blk_mq_rq_to_pdu(origin_rq);
	dev = dev_req->dev;

	blk_mq_start_request(origin_rq);
	path_queue = bdev_get_queue(dev->dev_block);
	clone_rq = blk_mq_alloc_request(path_queue, origin_rq->cmd_flags | REQ_NOMERGE, BLK_MQ_REQ_NOWAIT);
	if (IS_ERR(clone_rq)) {
		return BLK_STS_RESOURCE;
	}
	clone_rq->bio = clone_rq->biotail = NULL;
	clone_rq->cmd_flags |= REQ_FAILFAST_TRANSPORT;
	clone_rq->end_io = end_clone_request;
	clone_rq->end_io_data = dev_req;

	dev_req->clone_req = clone_rq;
	dev_req->dev = dev;
	dev_req->origin_req = origin_rq;
	dev_req->n_sectors = 0;
	dev_req->completed = 0;
	dev_req->error = 0;

	ret = blk_rq_prep_clone(clone_rq, origin_rq, &dev->dev_bio_set, GFP_ATOMIC, rq_bio_constructor, dev_req);
	if (ret) {
		return BLK_STS_RESOURCE;
	}

	ret = blk_insert_cloned_request(clone_rq);
	if (ret == BLK_STS_OK) {
		return BLK_STS_OK;
	} else {
		blk_rq_unprep_clone(clone_rq);
		blk_mq_cleanup_rq(clone_rq);
	}
	return ret;
}

static void device_complete(struct request *req)
{
	device_request_t *dev_req = NULL;
	dev_req = blk_mq_rq_to_pdu(req);
	blk_rq_unprep_clone(req);
	blk_mq_end_request(dev_req->clone_req, dev_req->error);
}

static int device_init_request(struct blk_mq_tag_set *set, struct request *rq, unsigned int hctx_idx, unsigned int numa_node)
{
	virtual_device_t *device = NULL;
	device_request_t *dev_req = NULL;
	device = set->driver_data;
	dev_req = blk_mq_rq_to_pdu(rq);
	dev_req->dev = device;
	return 0;
}

static const struct blk_mq_ops device_mq_ops = {
	.queue_rq = device_mq_queue_rq,
	.complete = device_complete,
	.init_request = device_init_request,
};

int init_device_queue(virtual_device_t *device)
{
	int err = -1;
	unsigned int front_pad = sizeof(clone_bio_t);
	bioset_init(&device->dev_bio_set, 256, front_pad, 0)
	device->dev_tag_set.ops = &device_mq_ops;
	device->dev_tag_set.queue_depth = 1024;
	device->dev_tag_set.numa_node = NUMA_NO_NODE;
	device->dev_tag_set.flags = BLK_MQ_F_SHOULD_MERGE | BLK_MQ_F_STACKING;
	device->dev_tag_set.nr_hw_queues = 1;
	device->dev_tag_set.driver_data = device;
	device->dev_tag_set.cmd_size = sizeof(device_request_t);	
	err = blk_mq_alloc_tag_set(&device->dev_tag_set);
	if (err) {
		return -1;
	}
	err = blk_mq_init_allocated_queue(&device->dev_tag_set, device->dev_queue);
	if (err) {
		blk_mq_free_tag_set(&device->dev_tag_set);
		return -1;
	}
	return 0;
}

void exit_device_queue(virtual_device_t *device)
{
    blk_mq_free_tag_set(&device->dev_tag_set);
}
